package websocket

import (
	"bytes"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"

	"time"
	"tuiwen-go/internal/api"
	"tuiwen-go/internal/config"
	"tuiwen-go/internal/db"
)

// taskImgHandler handles messages of type 4.
func taskImgHandler(c *Client, data json.RawMessage) {
	log.Printf("Client %s: Handling Image Task, Data: %s", c.ID, string(data))

	var taskData struct {
		ProjectID string `json:"project_id"`
		Id        string `json:"id"`
	} // Assuming the data contains project_id and project_name

	if err := json.Unmarshal(data, &taskData); err != nil {
		log.Printf("Client %s: Error unmarshalling image task data: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "Invalid image task data format"})
		return
	}

	if taskData.ProjectID == "" {
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "项目ID不能为空"})
		return
	}

	// Fetch text entries that need images generated
	sqlQuery := "SELECT id, prompt, image FROM text WHERE project_id = ? AND content IS NOT NULL"
	args := []interface{}{taskData.ProjectID}
	if taskData.Id != "" {
		sqlQuery += " AND id = ?"
		args = append(args, taskData.Id)
	} else {
		sqlQuery += " AND (image IS NULL OR image = '')"
	}
	rows, err := db.Query(sqlQuery, args...)
	if err != nil {
		log.Printf("Client %s: Error querying text for image task: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "查询待生成图片文本失败"})
		return
	}
	defer rows.Close()

	var textsToImage []struct {
		ID     int
		Prompt string
		Image  string // Add Image field
	}

	for rows.Next() {
		var id int
		var prompt sql.NullString
		var image sql.NullString // Scan the image column
		err := rows.Scan(&id, &prompt, &image)
		if err != nil {
			log.Printf("Client %s: Error scanning text row for image task: %v", c.ID, err)
			continue
		}

		if prompt.Valid {
			textsToImage = append(textsToImage, struct {
				ID     int
				Prompt string
				Image  string
			}{ID: id, Prompt: prompt.String, Image: image.String})
		}
	}

	if err = rows.Err(); err != nil {
		log.Printf("Client %s: Error iterating text rows for image task: %v", c.ID, err)
		sendMessage(c, api.WxJSONResponse{Code: -1, Msg: "处理待生成图片文本失败"})
		return
	}

	if len(textsToImage) == 0 {
		sendMessage(c, api.WxJSONResponse{Code: 0, Msg: "图片已全部生成"})
		return // No texts to generate images for
	}

	// Process each text entry to generate image
	for _, item := range textsToImage {
		// Placeholder for image generation logic
		// This requires calling an external service or a local model.
		// For now, let's use dummy values.
		// In a real implementation, integrate an image generation function.
		generatedImage, err := generateAndSaveImage(item.Prompt, taskData.ProjectID)
		if err != nil {
			log.Printf("Client %s: Error generating and saving image for text ID %d: %v", c.ID, item.ID, err)
			// Continue processing other items, but log the error
			sendMessage(c, api.WxJSONResponse{Code: 1, Msg: fmt.Sprintf("图片生成失败 (ID: %d): %v", item.ID, err)})
		} else {

			// Update the database with the generated image path
			sqlUpdate := "UPDATE text SET image = ? WHERE project_id = ? AND id = ?"
			_, err := db.Exec(sqlUpdate, generatedImage, taskData.ProjectID, item.ID)
			if err != nil {
				log.Printf("Client %s: Error updating image for text ID %d: %v", c.ID, item.ID, err)
				// Continue processing other items, but log the error
			} else {

				// 旧图片重命名，格式为：新图片名_旧图片名.png
				if item.Image != "" {
					oldImagePath := strings.ReplaceAll(item.Image, "/public", "./public")
					newImageFilename := path.Base(generatedImage)
					oldImageFilename := path.Base(oldImagePath)
					newOldImageName := strings.TrimSuffix(newImageFilename, path.Ext(newImageFilename)) + "_" + oldImageFilename
					oldImageDir := path.Dir(oldImagePath)
					newOldImagePath := path.Join(oldImageDir, newOldImageName)

					log.Printf("Client %s: Renaming old image file from %s to %s", c.ID, oldImagePath, newOldImagePath)
					err := os.Rename(oldImagePath, newOldImagePath)
					if err != nil {
						log.Printf("Client %s: Failed to rename old image file %s to %s: %v", c.ID, oldImagePath, newOldImagePath, err)
					}
				}
				// Send a message to the client with the updated image URL
				sendMessage(c, api.WxJSONResponse{
					Code: 1, // Or a specific code for image update
					Msg:  fmt.Sprintf("配图生成成功 (ID: %d)", item.ID),
					Data: map[string]interface{}{
						"id":         item.ID,
						"image_path": generatedImage,
					},
				})
			}
		}
	}

	// A final success message after attempting all updates (optional, as individual messages are sent)
	// sendMessage(c, api.WxJSONResponse{Code: 0, Msg: "图片生成任务完成"})
}

// generateAndSaveImage calls the Stable Diffusion API, saves the image, and returns the public path.
func generateAndSaveImage(prompt string, projectID string) (string, error) {
	log.Printf("Generating image for prompt: %s (Project ID: %s)", prompt, projectID)

	// Get Stable Diffusion URL from database
	stableDiffusionURL, ok := config.GetSetting("stable_diffusion_url")
	if !ok {
		log.Printf("获取 Stable Diffusion URL 设置失败: stable_diffusion_url not found in cache, trying database")
		stableDiffusionURL = "http://127.0.0.1:7860" // Fallback to default URL
	}

	if stableDiffusionURL == "" {
		// Fallback to default URL if the setting is empty in the database
		stableDiffusionURL = "http://127.0.0.1:7860"
	}
	widthStr, ok := config.GetSetting("image_width")
	if !ok {
		log.Printf("获取图片宽度设置失败: image_width not found in cache")
		widthStr = "800" // Fallback to default
	}
	heightStr, ok := config.GetSetting("image_height")
	if !ok {
		log.Printf("获取图片高度设置失败: image_height not found in cache")
		heightStr = "450" // Fallback to default
	}
	negativePrompt, ok := config.GetSetting("image_negative_prompt")
	if !ok {
		log.Printf("获取图片高度设置失败: negativePrompt not found in cache")
	}

	width, err := strconv.Atoi(widthStr)
	if err != nil {
		log.Printf("转换图片宽度失败: %v", err)
		width = 800
	}
	height, err := strconv.Atoi(heightStr)
	if err != nil {
		log.Printf("转换图片高度失败: %v", err)
		height = 450
	}
	image_sampler_name, ok := config.GetSetting("image_sampler_name")
	if !ok {
		log.Printf("获取图片高度设置失败: image_sampler_name not found in cache")
	}
	image_steps, ok := config.GetSetting("image_steps")
	if !ok {
		log.Printf("获取采样步数设置失败: image_steps not found in cache")
		image_steps = "30"
	}
	image_cfg_scale, ok := config.GetSetting("image_cfg_scale")
	if !ok {
		log.Printf("获取提示词引导强度设置失败: image_cfg_scale not found in cache")
		image_cfg_scale = "1"
	}
	image_scheduler, ok := config.GetSetting("image_scheduler")
	if !ok {
		log.Printf("获取提示词引导强度设置失败: image_cfg_scale not found in cache")
		image_scheduler = "Automatic"
	}
	// Prepare the request payload for Stable Diffusion API
	requestPayload := map[string]interface{}{
		"prompt": prompt,

		"width":           width,
		"height":          height,
		"negative_prompt": negativePrompt, // 负面提示词
		//	"seed":               -1,             // 随机种子，-1为随机
		//	"subseed":            -1,             // 次级种子
		//	"subseed_strength":   0.3,            // 次级种子影响强度（0-1）
		//	"seed_resize_from_h": -1,             // 高度方向上调整种子大小的来源
		//	"seed_resize_from_w": -1,             // 宽度方向上调整种子大小的来源
		"sampler_name": image_sampler_name, // 采样器名称
		"scheduler":    image_scheduler,    // 调度器类型
		//	"batch_size":         1,              // 单次生成的图像数量
		"n_iter":              1,               // 迭代次数
		"steps":               image_steps,     // 采样步数
		"cfg_scale":           image_cfg_scale, // 提示词引导强度
		"distilled_cfg_scale": 3.5,             // 蒸馏模型的CFG缩放
	}

	jsonPayload, err := json.Marshal(requestPayload)
	if err != nil {
		return "", fmt.Errorf("error marshalling request payload: %v", err)
	}

	// Make the HTTP POST request to the Stable Diffusion API
	resp, err := http.Post(stableDiffusionURL+"/sdapi/v1/txt2img", "application/json", bytes.NewBuffer(jsonPayload))
	if err != nil {
		return "", fmt.Errorf("error making HTTP request to Stable Diffusion API: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("Stable Diffusion API returned non-OK status code: %d, body: %s", resp.StatusCode, string(bodyBytes))
	}

	// Read the response body
	bodyBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("error reading response body from Stable Diffusion API: %v", err)
	}

	// Unmarshal the response body
	var response struct {
		Images []string `json:"images"`
	}
	if err := json.Unmarshal(bodyBytes, &response); err != nil {
		return "", fmt.Errorf("error unmarshalling response body from Stable Diffusion API: %v", err)
	}

	if len(response.Images) == 0 {
		return "", fmt.Errorf("no images returned from Stable Diffusion API")
	}

	// Decode the base64 image data
	base64Data := response.Images[0]
	// Remove potential data URL prefix
	if strings.HasPrefix(base64Data, "data:image/") {
		parts := strings.SplitN(base64Data, ",", 2)
		if len(parts) == 2 {
			base64Data = parts[1]
		}
	}

	imageData, err := base64.StdEncoding.DecodeString(base64Data)
	if err != nil {
		return "", fmt.Errorf("error decoding base64 image data: %v", err)
	}

	dir := fmt.Sprintf("public/images/%s", projectID)
	// Ensure the directory exists relative to the executable
	executableDir, err := os.Executable()
	if err != nil {
		return "", fmt.Errorf("error getting executable directory: %v", err)
	}
	executablePath := path.Dir(executableDir)

	saveDir := path.Join(executablePath, dir)

	if _, err := os.Stat(saveDir); os.IsNotExist(err) {
		err = os.MkdirAll(saveDir, 0755)
		if err != nil {
			return "", fmt.Errorf("error creating directory %s: %v", saveDir, err)
		}
	}

	filename := fmt.Sprintf("%d.png", time.Now().UnixNano())
	savePath := path.Join(saveDir, filename)

	// Save the image file
	err = os.WriteFile(savePath, imageData, 0644)
	if err != nil {
		return "", fmt.Errorf("error saving image file %s: %v", savePath, err)
	}

	// Return the public path relative to the 'public' directory
	// This path will be used by the frontend to access the image via the static file server.
	publicPath := path.Join("images", projectID, filename)
	// Use forward slashes for URL path regardless of OS
	publicPath = strings.ReplaceAll(publicPath, "\\", "/")

	log.Printf("Image saved to %s, public path: %s", savePath, publicPath)

	return "/public/" + publicPath, nil
}
